import datetime
from typing import List, Dict

from sqlalchemy import Column, DateTime, Date, JSON, String, TEXT, Index, func, and_, BIGINT
from sqlalchemy.orm import Session

from common.database import Base
from common.database import SessionMixin
from core.logger import logger


class ImpalaQueryInfo(Base, SessionMixin):
    __tablename__ = 'bigdata_impala_query_info'
    query_id = Column(String(128), nullable=False, comment="查询ID")
    user = Column(String(128), nullable=False, comment="用户")
    statement = Column(TEXT, nullable=True, comment="查询语句")
    duration_millis = Column(BIGINT, default=0, nullable=False, comment="执行时长")
    query_state = Column(String(20), nullable=True, comment="查询状态")
    query_status = Column(TEXT, nullable=True, comment="状态描述")
    memory_aggregate_peak = Column(BIGINT, default=0, nullable=False, comment="节点内存使用峰值")
    estimated_per_node_peak_memory = Column(BIGINT, default=0, nullable=False, comment="节点内存估算峰值")
    thread_cpu_time = Column(BIGINT, default=0, nullable=False, comment="所有查询线程使用CPU时间总和")
    oom = Column(String(20), nullable=True, comment="内存溢出状态")
    start_time = Column(DateTime, nullable=False, comment="开始时间")
    end_time = Column(DateTime, nullable=True, comment="结束时间")
    query_response = Column(JSON, nullable=False, comment="查询结果")
    create_date = Column(Date, nullable=False, comment="创建日期")
    create_time = Column(DateTime, default=datetime.datetime.now, nullable=False, comment="创建时间")

    __table_args__ = (
        Index('idx_bigdata_impala_query_info_uuid', 'uuid', unique=True),
        Index('idx_bigdata_impala_query_info_query_id', 'query_id'),
        Index('idx_bigdata_impala_query_info_create_date', 'create_date'),
        Index('idx_bigdata_impala_query_info_create_time', 'create_time'),
        Index('idx_bigdata_impala_query_info_start_time', 'start_time'),
        Index('idx_bigdata_impala_query_info_end_time', 'end_time'),
        Index('idx_bigdata_impala_create_date_duration', 'create_date', 'duration_millis'),
        Index('idx_bigdata_impala_create_date_aggregate', 'create_date', 'memory_aggregate_peak'),
        Index('idx_bigdata_impala_create_date_estimated', 'create_date', 'estimated_per_node_peak_memory'),
        Index('idx_bigdata_impala_create_date_thread', 'create_date', 'thread_cpu_time'),
        Index('idx_bigdata_impala_create_time_duration', 'start_time', 'duration_millis'),
        Index('idx_bigdata_impala_create_time_aggregate', 'start_time', 'memory_aggregate_peak'),
        Index('idx_bigdata_impala_create_time_estimated', 'start_time', 'estimated_per_node_peak_memory'),
        Index('idx_bigdata_impala_create_time_thread', 'start_time', 'thread_cpu_time'),
        {
            'mysql_engine': 'innodb',
            'mysql_charset': 'utf8mb4',
            'mysql_collate': 'utf8mb4_unicode_ci',
            'extend_existing': True,
            'mysql_row_format': 'dynamic',
        }
    )

    def __repr__(self):
        return "<user='{} ,query_id='{} '>".format(self.user, self.query_id)

    @classmethod
    def query_update(cls, db: Session, data: dict):
        try:
            result = db.query(cls).filter_by(query_id=data['query_id']).update(data)
            db.commit()
        except Exception as e:
            db.rollback()
            logger.error(f'update error => {e}')
            return None
        return result

    @classmethod
    def get_response_by_query_id(cls, db: Session, query_id: str) -> Dict:
        return db.query(cls).with_entities(cls.statement).filter(cls.query_id == query_id).to_dict_first()

    @classmethod
    def get_statement_by_range_datetime(cls, db: Session, start_datatime: str, end_datatime: str, statement="compute stats%") -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.statement, cls.duration_millis, cls.query_response, cls.start_time,
                                           cls.end_time).filter(
            and_(cls.start_time.between(start_datatime, end_datatime), cls.statement.like(statement))).order_by(
            cls.thread_cpu_time.desc()).to_dict()

    @classmethod
    def get_duration_millis_top_by_date(cls, db: Session, date: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.create_date == date).order_by(cls.duration_millis.desc()).limit(top).to_dict()

    @classmethod
    def get_duration_millis_top_by_range_datetime(cls, db: Session, start_datatime: str, end_datatime: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.start_time.between(start_datatime, end_datatime)).order_by(cls.duration_millis.desc()).limit(top).to_dict()

    @classmethod
    def get_memory_aggregate_top_by_date(cls, db: Session, date: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.create_date == date).order_by(cls.memory_aggregate_peak.desc()).limit(top).to_dict()

    @classmethod
    def get_memory_aggregate_top_by_range_datetime(cls, db: Session, start_datatime: str, end_datatime: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.start_time.between(start_datatime, end_datatime)).order_by(cls.memory_aggregate_peak.desc()).limit(top).to_dict()

    @classmethod
    def get_memory_estimated_top_by_date(cls, db: Session, date: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.create_date == date).order_by(cls.estimated_per_node_peak_memory.desc()).limit(top).to_dict()

    @classmethod
    def get_memory_estimated_top_by_range_datetime(cls, db: Session, start_datatime: str, end_datatime: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.start_time.between(start_datatime, end_datatime)).order_by(cls.estimated_per_node_peak_memory.desc()).limit(
            top).to_dict()

    @classmethod
    def get_tread_cpu_time_top_by_date(cls, db: Session, date: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.create_date == date).order_by(cls.thread_cpu_time.desc()).limit(top).to_dict()

    @classmethod
    def get_tread_cpu_time_top_by_range_datetime(cls, db: Session, start_datatime: str, end_datatime: str, top: int = 10) -> List:
        return db.query(cls).with_entities(cls.query_id, cls.user, cls.duration_millis, cls.memory_aggregate_peak,
                                           cls.estimated_per_node_peak_memory, cls.thread_cpu_time, cls.start_time, cls.end_time) \
            .filter(cls.start_time.between(start_datatime, end_datatime)).order_by(cls.thread_cpu_time.desc()).limit(top).to_dict()


if __name__ == '__main__':
    from common.database import SessionLocal

    session = SessionLocal()
    impala = ImpalaQueryInfo()
    # print(impala.get_duration_millis_top_by_date("2023-02-14", session))
    # print(impala.get_memory_aggregate_top_by_date("2023-02-14", session))
    impala.get_statement_by_range_datetime(session, "2024-01-12 00:00:00", "2024-01-12 23:59:59")
    # print(impala.get_response_by_query_id(session, "1449d27d2db5ea83:655f35e500000000"))
    # print(impala.get_memory_estimated_top_by_date("2023-02-14", session))
    # print(kudu_table_info.get_latest_time_range(session))
